from rag_qa.core.prompts import RAGPrompts
import time
from base import logger, Config
from rag_qa.core.query_classifier import QueryClassifier  # 查询分类器
from rag_qa.core.strategy_selector import StrategySelector  # 策略选择器

config = Config()


# 定义 RAGSystem 类，封装 RAG 系统的核心逻辑
class RAGSystem:
	# 初始化方法，设置 RAG 系统的基本参数
	def __init__(self, vector_store, llm):
		# 设置向量数据库对象
		self.vector_store = vector_store
		# 设置 LLM 调用函数
		self.llm = llm
		# 获取 RAG 提示模板
		self.rag_prompts = RAGPrompts.rag_prompt()
		# 初始化查询分类器
		self.query_classifier = QueryClassifier(
			model_path=r'F:\Study\Project\EduRAGProject\rag_qa\core\bert_query_classifier'
		)
		# 初始化策略选择器
		self.strategy_selector = StrategySelector()
		# 最大 4096 字符
		self.max_prompt_length = 4096

	# 定义私有方法，使用假设文档进行检索（HyDE）
	def _retrieve_with_hyde(self, query):
		logger.info(f"使用 HyDE 策略进行检索 (查询：'{query}')")
		# 获取假设向量生成的 Prompt 模版
		hyde_prompt_template = RAGPrompts.hyde_prompt()  # 使用 template 后缀区分
